import {
    DynamicModule,
    Global,
    Inject,
    Module,
    OnApplicationShutdown,
    Provider,
    Type,
} from '@nestjs/common';
import { ModuleRef } from '@nestjs/core';
import * as mongoose from 'mongoose';
import { ConnectOptions, Connection } from 'mongoose';
import { defer, lastValueFrom } from 'rxjs';
import { catchError } from 'rxjs/operators';
import { getConnectionToken, handleRetry } from './mongoose.utils';
import {
    MongooseModuleAsyncOptions,
    MongooseModuleFactoryOptions,
    MongooseModuleOptions,
    MongooseOptionsFactory,
} from './mongoose-options.interface';
import {
    MONGOOSE_CONNECTION_NAME,
    MONGOOSE_MODULE_OPTIONS,
} from './mongoose.constants';

@Global()
@Module({})
export class MongooseCoreModule implements OnApplicationShutdown {

    private static connections: Record<string, any> = {} 

    constructor(
        @Inject(MONGOOSE_CONNECTION_NAME) private readonly connectionName: string,
        private readonly moduleRef: ModuleRef,
    ) { }

    static forRoot(
        uri: string,
        options: MongooseModuleOptions = {},
    ): DynamicModule {
        const {
            retryAttempts,
            retryDelay,
            connectionName,
            connectionFactory,
            connectionErrorFactory,
            lazyConnection,
            onConnectionCreate,
            verboseRetryLog,
            ...mongooseOptions
        } = options;

        const mongooseConnectionFactory =
            connectionFactory || ((connection) => connection);

        const mongooseConnectionError =
            connectionErrorFactory || ((error) => error);

        const mongooseConnectionName = getConnectionToken(connectionName);

        const mongooseConnectionNameProvider = {
            provide: MONGOOSE_CONNECTION_NAME,
            useValue: mongooseConnectionName,
        };

        const connectionProvider = {
            provide: mongooseConnectionName,
            useFactory: async (): Promise<any> =>
                await lastValueFrom(
                    defer(async () =>
                        mongooseConnectionFactory(
                            await this.createMongooseConnection(uri, mongooseOptions, {
                                lazyConnection,
                                onConnectionCreate,
                            }),
                            mongooseConnectionName,
                        ),
                    ).pipe(
                        handleRetry(retryAttempts, retryDelay, verboseRetryLog),
                        catchError((error) => {
                            throw mongooseConnectionError(error);
                        }),
                    ),
                ),
        };
        return {
            module: MongooseCoreModule,
            providers: [connectionProvider, mongooseConnectionNameProvider],
            exports: [connectionProvider],
        };
    }

    static forRootAsync(options: MongooseModuleAsyncOptions): DynamicModule {
        const mongooseConnectionName = getConnectionToken(options.connectionName);

        const mongooseConnectionNameProvider = {
            provide: MONGOOSE_CONNECTION_NAME,
            useValue: mongooseConnectionName,
        };

        const connectionProvider = {
            provide: mongooseConnectionName,
            useFactory: async (
                mongooseModuleOptions: MongooseModuleFactoryOptions,
            ): Promise<any> => {
                if(!mongooseModuleOptions) return;

                const {
                    retryAttempts,
                    retryDelay,
                    uri,
                    connectionFactory,
                    connectionErrorFactory,
                    lazyConnection,
                    onConnectionCreate,
                    verboseRetryLog,
                    ...mongooseOptions
                } = mongooseModuleOptions;

                const mongooseConnectionFactory =
                    connectionFactory || ((connection) => connection);

                const mongooseConnectionError =
                    connectionErrorFactory || ((error) => error);

                return await lastValueFrom(
                    defer(async () =>
                        mongooseConnectionFactory(
                            await this.createMongooseConnection(
                                uri as string,
                                mongooseOptions,
                                { lazyConnection, onConnectionCreate },
                            ),
                            mongooseConnectionName,
                        ),
                    ).pipe(
                        handleRetry(retryAttempts, retryDelay, verboseRetryLog),
                        catchError((error) => {
                            throw mongooseConnectionError(error);
                        }),
                    ),
                );
            },
            inject: [MONGOOSE_MODULE_OPTIONS],
        };
        const asyncProviders = this.createAsyncProviders(options);
        return {
            module: MongooseCoreModule,
            imports: options.imports,
            providers: [
                ...asyncProviders,
                connectionProvider,
                mongooseConnectionNameProvider,
            ],
            exports: [connectionProvider],
        };
    }

    private static createAsyncProviders(
        options: MongooseModuleAsyncOptions,
    ): Provider[] {
        if (options.useExisting || options.useFactory) {
            return [this.createAsyncOptionsProvider(options)];
        }
        const useClass = options.useClass as Type<MongooseOptionsFactory>;
        return [
            this.createAsyncOptionsProvider(options),
            {
                provide: useClass,
                useClass,
            },
        ];
    }

    private static createAsyncOptionsProvider(
        options: MongooseModuleAsyncOptions,
    ): Provider {
        if (options.useFactory) {
            return {
                provide: MONGOOSE_MODULE_OPTIONS,
                useFactory: options.useFactory,
                inject: options.inject || [],
            };
        }
        // `as Type<MongooseOptionsFactory>` is a workaround for microsoft/TypeScript#31603
        const inject = [
            (options.useClass || options.useExisting) as Type<MongooseOptionsFactory>,
        ];
        return {
            provide: MONGOOSE_MODULE_OPTIONS,
            useFactory: async (optionsFactory: MongooseOptionsFactory) =>
                await optionsFactory.createMongooseOptions(),
            inject,
        };
    }

    private static async createMongooseConnection(
        uri: string,
        mongooseOptions: ConnectOptions,
        factoryOptions: {
            lazyConnection?: boolean;
            onConnectionCreate?: MongooseModuleOptions['onConnectionCreate'];
        },
    ): Promise<Connection> {
        if(this.connections[uri]) {
            return this.connections[uri]
        }

        const connection = mongoose.createConnection(uri, mongooseOptions);
        this.connections[uri] = connection

        if (factoryOptions?.lazyConnection) {
            return connection;
        }

        factoryOptions?.onConnectionCreate?.(connection);

        return connection.asPromise();
    }

    async onApplicationShutdown() {
        const connection = this.moduleRef.get<any>(this.connectionName);
        if (connection) {
            await connection.close();
        }

        if (MongooseCoreModule.connections && Object.keys(MongooseCoreModule.connections).length > 0) {
            for (const key of Object.keys(MongooseCoreModule.connections)) {
                const client = MongooseCoreModule.connections[key]
                if (client && typeof client.close == 'function') {
                    client.close()
                }
            }
        }
    }
}